package com.mygdx.game.dl4j;

import com.badlogic.gdx.Application;
import com.badlogic.gdx.backends.lwjgl.LwjglApplication;
import com.badlogic.gdx.backends.lwjgl.LwjglApplicationConfiguration;
import com.re6g3y.dl4j.Setting;
import com.re6g3y.dl4j.core.Client;
import com.re6g3y.dl4j.core.GameCore;
import org.deeplearning4j.gym.StepReply;
import org.deeplearning4j.rl4j.mdp.MDP;
import org.deeplearning4j.rl4j.space.ActionSpace;
import org.deeplearning4j.rl4j.space.DiscreteSpace;
import org.deeplearning4j.rl4j.space.Encodable;
import org.deeplearning4j.rl4j.space.ObservationSpace;
import org.nd4j.linalg.api.ndarray.INDArray;

public class CarMDP<O extends Encodable, A, AS extends ActionSpace<A>> implements MDP<O, A, AS> {

    Application app;
    Client<A,StepReply<O>> client;
    public CarMDP(){
        LwjglApplicationConfiguration config = new LwjglApplicationConfiguration();
        config.width=(int)(1920/1.5f);
        config.foregroundFPS=600000;
        config.backgroundFPS=600000;
        config.vSyncEnabled=true;
        config.height=(int)(1080/1.5f);
//		new LwjglApplication(new MyGdxGame(), config);
//		new LwjglApplication(new ball(), config);
        client=new Client<A,StepReply<O>>();
        app=new LwjglApplication(new GameCore(client), config);
    }

    private ActionSpace<Integer> as= new DiscreteSpace(Setting.ActionNumber);
    private ObservationSpace<O> observationSpace=new ObservationSpace<O>() {
        @Override
        public String getName() {
            return "CarMDP-ObservationSpace";
        }

        @Override
        public int[] getShape() {
            return new int[]{Setting.SensorNumber};
        }

        @Override
        public INDArray getLow() {
            return null;
        }

        @Override
        public INDArray getHigh() {
            return null;
        }
    };

    @Override
    public ObservationSpace<O> getObservationSpace() {
        return observationSpace;
    }

    @Override
    public AS getActionSpace() {
        return (AS)as;
    }

    @Override
    public O reset() {
        A act=(A)(Integer)1;
        client.reset();
        client.setAction(act);
        StepReply s=client.getStepReply();
        done=s.isDone();
        return (O)s.getObservation();
    }

    @Override
    public void close() {
        client.exit();
        app.exit();
        System.exit(0);
    }

    boolean done=true;
    @Override
    public StepReply<O> step(A act) {
        client.setAction(act);
        StepReply s=client.getStepReply();
        done=s.isDone();
        return s;
    }

    @Override
    public boolean isDone() {
        return done;
    }

    @Override
    public MDP<O, A, AS> newInstance() {
        return new CarMDP<O, A, AS>();
    }
}
